import zipfile
import xml.etree.ElementTree as ET
from pathlib import Path
from docx.table import Table
import os
import os.path as osp
from tabel_and_images.read_image import parse_image_llava
p = '/home/FAST_DATA_MIRROR/Langchain-Chatchat-master/tabel/test2.docx'
tard = Path(p).parent.joinpath(Path(p).stem)
# 解压docx文件以访问其内部XML结构
with zipfile.ZipFile(p, 'r') as docx:
    docx.extractall(tard)
# 图像保存文件夹
images_folder = f'{tard}/extracted_images'
Path(images_folder).mkdir(parents=True, exist_ok=True)

# 加载文档主体的XML
tree = ET.parse(f'{tard}/word/document.xml')
root = tree.getroot()

# 解析document.xml.rels文件获取关系标识符(即rId)到图像路径的映射
rels_path = f'{tard}/word/_rels/document.xml.rels'
tree_rels = ET.parse(rels_path)
root_rels = tree_rels.getroot()
# 建立rId到图像文件名的映射
rId_to_image = {
    rel.attrib['Id']: rel.attrib['Target']
    for rel in root_rels.findall('.//{http://schemas.openxmlformats.org/package/2006/relationships}Relationship')
    if rel.attrib['Type'].endswith('/image')
}
# 定义命名空间
namespaces = {
    'w': 'http://schemas.openxmlformats.org/wordprocessingml/2006/main',
    'a': 'http://schemas.openxmlformats.org/drawingml/2006/main',
    'r': 'http://schemas.openxmlformats.org/officeDocument/2006/relationships',
    'wp': 'http://schemas.openxmlformats.org/drawingml/2006/wordprocessingDrawing',
}
# 遍历文档中的段落和表格
for child in root.iter():
    if child.tag.endswith('tbl'):
        print("Table found")
        # 遍历表格中的所有行和单元格
        rows = []
        for row in child.findall('.//w:tr', namespaces):
            for cell in row.findall('.//w:tc', namespaces):
                # 检查单元格内的文本
                texts = [node.text for node in cell.findall('.//w:t', namespaces)]
                texts = ''.join(texts)
                rows.append(texts)
                # 检查图像
                drawing = cell.find('.//w:drawing', namespaces)
                if drawing is not None:
                    for drawing in drawings:
                            # 找到了图片
                            blips = drawing.findall('.//a:blip', namespaces)
                            for blip in blips:
                                rId = blip.attrib['{http://schemas.openxmlformats.org/officeDocument/2006/relationships}embed']
                                if rId in rId_to_image:
                                    image_path = rId_to_image[rId]
                                    image_path = f"{tard}/word/{rId_to_image[rId]}"
                                    output = parse_image_llava(image_path=image_path)
            print(rows) 
                    
    elif child.tag.endswith('p'):
        # 这是段落的逻辑...
        # 遍历段落查找文本
        texts = [node.text for node in child.findall('.//w:t', namespaces)]
        paragraph_text = ''.join(texts)
        if len(paragraph_text)<=0:
            continue
        print('Paragraph:', paragraph_text)
        
        # 查找段落中的图像
        drawings = child.findall('.//w:drawing', namespaces)
        for drawing in drawings:
            # 找到了图片
            for drawing in drawings:
                    # 找到了图片
                    blips = drawing.findall('.//a:blip', namespaces)
                    for blip in blips:
                        rId = blip.attrib['{http://schemas.openxmlformats.org/officeDocument/2006/relationships}embed']
                        if rId in rId_to_image:
                            image_path = f"{tard}/word/{rId_to_image[rId]}"
                            assert osp.exists(image_path)
                            output = parse_image_llava(image_path=image_path)
